#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.0
# @author : cd
# @Email : 19688513@qq.com
# @Project : new-horizons-engine
# @File : test.py
# @Software: PyCharm
# @time: 2025/6/6 17:17
# @description : 
"""

import urllib3
import logging
import akshare as ak
import pandas as pd
import re

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 禁用不安全的请求警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)




class AKShareDataFetcher:
    """
    统一封装东方财富、新浪、腾讯股票历史行情数据接口
    输出统一格式：日期(date), 开盘(open), 收盘(close), 最高(high), 最低(low), 成交量(volume, 单位:股), 成交额(amount, 单位:元)
    """

    @staticmethod
    def _standardize_output(df, volume_unit='手'):
        """
        标准化输出格式和单位
        :param df: 原始DataFrame
        :param volume_unit: 原始成交量单位 ('手' 或 '股')
        :return: 标准化后的DataFrame
        """
        # 列名标准化
        col_map = {
            '日期': 'date', '股票代码': 'symbol', '开盘': 'open', '收盘': 'close',
            '最高': 'high', '最低': 'low', '成交量': 'volume', '成交额': 'amount'
        }
        df = df.rename(columns=col_map)

        # 保留需要的列
        keep_cols = ['date', 'open', 'close', 'high', 'low', 'volume', 'amount']
        df = df[[col for col in keep_cols if col in df.columns]]

        # 单位转换
        if 'volume' in df.columns:
            if volume_unit == '手':
                df['volume'] = df['volume'] * 100  # 手 -> 股
            elif volume_unit == '股':
                pass  # 已经是股，无需转换
            else:
                raise ValueError(f"未知成交量单位: {volume_unit}")

        # 确保数据类型正确
        num_cols = ['open', 'close', 'high', 'low', 'volume', 'amount']
        for col in num_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')

        return df

    @staticmethod
    def _get_market_prefix(symbol):
        """获取股票代码的市场前缀"""
        if symbol.startswith(('sh', 'sz', 'bj')):
            return symbol[:2].lower()
        if symbol.startswith(('6', '9')):
            return 'sh'
        elif symbol.startswith(('0', '3')):
            return 'sz'
        elif symbol.startswith('4') or symbol.startswith('8'):
            return 'bj'
        else:
            raise ValueError(f"无法识别的股票代码: {symbol}")

    def fetch_em(self, symbol, period='daily', start_date='19900101', end_date='20500101', adjust=''):
        """
        东方财富接口
        :param symbol: 股票代码 (带或不带市场标识)
        :param period: 周期 ('daily', 'weekly', 'monthly')
        :param start_date: 开始日期 (YYYYMMDD)
        :param end_date: 结束日期 (YYYYMMDD)
        :param adjust: 复权类型 ('', 'qfq', 'hfq')
        :return: 标准化DataFrame
        """
        # 处理股票代码格式
        clean_symbol = re.sub(r'\D', '', symbol)[-6:]

        # 获取数据
        df = ak.stock_zh_a_hist(
            symbol=clean_symbol,
            period=period,
            start_date=start_date,
            end_date=end_date,
            adjust=adjust
        )

        # 标准化输出
        return self._standardize_output(df, volume_unit='手')

    def fetch_sina(self, symbol, start_date='19900101', end_date='20500101', adjust=''):
        """
        新浪财经接口
        :param symbol: 股票代码 (带或不带市场标识)
        :param start_date: 开始日期 (YYYYMMDD)
        :param end_date: 结束日期 (YYYYMMDD)
        :param adjust: 复权类型 ('', 'qfq', 'hfq')
        :return: 标准化DataFrame
        """
        # 处理股票代码格式
        prefix = self._get_market_prefix(symbol)
        clean_symbol = prefix + symbol[-6:]

        # 获取数据
        df = ak.stock_zh_a_daily(
            symbol=clean_symbol,
            start_date=start_date,
            end_date=end_date,
            adjust=adjust
        )

        # 标准化输出
        return self._standardize_output(df, volume_unit='股')

    def fetch_qq(self, symbol, start_date='19900101', end_date='20500101', adjust=''):
        """
        腾讯财经接口
        :param symbol: 股票代码 (带或不带市场标识)
        :param start_date: 开始日期 (YYYYMMDD)
        :param end_date: 结束日期 (YYYYMMDD)
        :param adjust: 复权类型 ('', 'qfq', 'hfq')
        :return: 标准化DataFrame
        """
        # 处理股票代码格式
        prefix = self._get_market_prefix(symbol)
        clean_symbol = prefix + symbol[-6:]

        # 获取数据
        df = ak.stock_zh_a_hist_tx(
            symbol=clean_symbol,
            start_date=start_date,
            end_date=end_date,
            adjust=adjust
        )

        # 腾讯接口没有成交额，添加空列
        if 'amount' not in df.columns:
            df['amount'] = None

        # 标准化输出
        return self._standardize_output(df, volume_unit='手')

    def get_data(self, source='em', **kwargs):
        """
        统一获取数据的方法
        :param source: 数据源 ('em', 'sina', 'qq')
        :param kwargs: 对应接口的参数
        :return: 标准化DataFrame
        """
        source_methods = {
            'em': self.fetch_em,
            'sina': self.fetch_sina,
            'qq': self.fetch_qq
        }

        if source not in source_methods:
            raise ValueError(f"无效的数据源: {source}，可选值: 'em', 'sina', 'qq'")

        return source_methods[source](**kwargs)


# 使用示例
if __name__ == "__main__":
    fetcher = AKShareDataFetcher()

    # 示例1: 获取东方财富数据 (不复权)
    print("东方财富数据示例:")
    df_em = fetcher.get_data(
        source='em',
        symbol='600000',
        period='daily',
        start_date='20230101',
        end_date='20230630',
        adjust='qfq'
    )
    print(df_em.head())

    # 示例2: 获取新浪财经数据 (前复权)
    print("\n新浪财经数据示例:")
    df_sina = fetcher.get_data(
        source='sina',
        symbol='sh600000',
        start_date='20230101',
        end_date='20230630',
        adjust='qfq'
    )
    print(df_sina.head())

    # 示例3: 获取腾讯财经数据 (后复权)
    print("\n腾讯财经数据示例:")
    df_qq = fetcher.get_data(
        source='qq',
        symbol='600000',
        start_date='20230101',
        end_date='20230630',
        adjust='qfq'
    )
    print(df_qq.head())